{"cells":[{"cell_type":"markdown","source":["Copyright 2022 Google LLC. SPDX-License-Identifier: Apache-2.0\n","\n","Licensed under the Apache License, Version 2.0 (the \"License\"); you may not use this file except in compliance with the License. You may obtain a copy of the License at\n","\n","https://www.apache.org/licenses/LICENSE-2.0\n","\n","Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License."],"metadata":{"id":"vyhthcZwGX9w"}},{"cell_type":"markdown","source":["# Experiment: HumanEval Benchmark\n","\n","This notebook is a part of the open-source code release associated with the paper:\n","\n","[Code as Policies: Language Model Programs for Embodied Control](https://code-as-policies.github.io/)\n","\n","This notebook gives the results corresponding to Table III in the paper which evaluates different code-gen approaches on the [HumanEval benchmark](https://github.com/openai/human-eval)\n","\n","1) Please obtain an OpenAI API Key here:\n","https://openai.com/blog/openai-api/\n","\n","2) Gain Codex access by joining the waitlist here:\n","https://openai.com/blog/openai-codex/\n","\n","Once you have Codex access you can use `code-davinci-002` as the `model_name`. Using the GPT-3 model (`text-dainvci-002`) is also ok, but performance won't be as good (there will be more code logic errors).\n","\n","3) Please also specify a location in your Google Drive on which the results will be stored.\n","\n","Note due to current rate limiting of the Codex API, this entire notebook may take 20+ hours to finish."],"metadata":{"id":"4rnus7D4GY60"}},{"cell_type":"code","source":["openai_api_key = 'YOUR KEY HERE'\n","model_name = 'code-davinci-002' # 'text-davinci-002'\n","google_drive_folder = 'drive/MyDrive/...'"],"metadata":{"id":"_0VeqEuzG97K"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"RWsZTVWhMH43"},"source":["# HumanEval CodeGen Benchmark\n","\n","From [HumanEval](https://github.com/openai/human-eval)"]},{"cell_type":"code","execution_count":null,"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"elapsed":30495,"status":"ok","timestamp":1662041664535,"user":{"displayName":"Jacky Liang","userId":"05524594537381871813"},"user_tz":240},"id":"ib2znWjv1XXA","outputId":"e6ca301a-395a-459f-c44a-3d6aa39fb687"},"outputs":[{"output_type":"stream","name":"stdout","text":["Mounted at /content/drive\n"]}],"source":["from google.colab import drive\n","drive.mount('/content/drive')"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"YGLNM9CE1kA0"},"outputs":[],"source":["from pathlib import Path\n","results_path = Path(google_drive_folder)"]},{"cell_type":"markdown","metadata":{"id":"QrSUzlIxOytF"},"source":["# Install HumanEval benchmark"]},{"cell_type":"code","execution_count":null,"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"elapsed":13378,"status":"ok","timestamp":1662041677911,"user":{"displayName":"Jacky Liang","userId":"05524594537381871813"},"user_tz":240},"id":"Q1rJ1hKxMJGe","outputId":"3eea7edd-0b11-4f81-b232-5e8654dd8069"},"outputs":[{"output_type":"stream","name":"stdout","text":["Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n","Collecting git+https://github.com/openai/human-eval.git\n","  Cloning https://github.com/openai/human-eval.git to /tmp/pip-req-build-842wcofg\n","  Running command git clone -q https://github.com/openai/human-eval.git /tmp/pip-req-build-842wcofg\n","Requirement already satisfied: tqdm in /usr/local/lib/python3.7/dist-packages (from human-eval==1.0) (4.64.0)\n","Collecting fire\n","  Downloading fire-0.4.0.tar.gz (87 kB)\n","\u001b[K     |████████████████████████████████| 87 kB 3.2 MB/s \n","\u001b[?25hRequirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from human-eval==1.0) (1.21.6)\n","Requirement already satisfied: six in /usr/local/lib/python3.7/dist-packages (from fire->human-eval==1.0) (1.15.0)\n","Requirement already satisfied: termcolor in /usr/local/lib/python3.7/dist-packages (from fire->human-eval==1.0) (1.1.0)\n","Building wheels for collected packages: human-eval, fire\n","  Building wheel for human-eval (setup.py) ... \u001b[?25l\u001b[?25hdone\n","  Created wheel for human-eval: filename=human_eval-1.0-py3-none-any.whl size=7446 sha256=ad3ee05f24b8af20f3e44126c16576373922992534af96274bf4455fb7db8a1b\n","  Stored in directory: /tmp/pip-ephem-wheel-cache-l44fg2ie/wheels/10/c6/41/a3d3cf28a68aa72be379d082afbafcc713353941c175b69b2d\n","  Building wheel for fire (setup.py) ... \u001b[?25l\u001b[?25hdone\n","  Created wheel for fire: filename=fire-0.4.0-py2.py3-none-any.whl size=115942 sha256=ba56f040c0ed9162515499f9d92f6994bde51e40a1391f1b6be41e040660e01b\n","  Stored in directory: /root/.cache/pip/wheels/8a/67/fb/2e8a12fa16661b9d5af1f654bd199366799740a85c64981226\n","Successfully built human-eval fire\n","Installing collected packages: fire, human-eval\n","\u001b[31mERROR: For req: human-eval==1.0. Invalid script entry point: <ExportEntry evaluate_functional_correctness = human_eval.evaluate_functional_correctness:None []> - A callable suffix is required. Cf https://packaging.python.org/specifications/entry-points/#use-for-scripts for more information.\u001b[0m\n","Cloning into 'human-eval'...\n","remote: Enumerating objects: 29, done.\u001b[K\n","remote: Counting objects: 100% (11/11), done.\u001b[K\n","remote: Compressing objects: 100% (10/10), done.\u001b[K\n","remote: Total 29 (delta 4), reused 1 (delta 1), pack-reused 18\u001b[K\n","Unpacking objects: 100% (29/29), done.\n"]}],"source":["!pip install git+https://github.com/openai/human-eval.git\n","\n","# Their package also needs some data not installed by pip.\n","# Get this straight from git cloning the repo,\n","# and hacking some paths.\n","!git clone https://github.com/openai/human-eval\n","!cp -r human-eval/data /usr/local/lib/python3.7/dist-packages"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"DSMKcSI7-1nr"},"outputs":[],"source":["# To make the evaluation actually run,\n","# we have to uncomment a line that they left commented as a safety.\n","!sed -i '58s/# //' /usr/local/lib/python3.7/dist-packages/human_eval/execution.py"]},{"cell_type":"markdown","metadata":{"id":"6O-5aNIPO2OQ"},"source":["# Install Codex LLM"]},{"cell_type":"code","execution_count":null,"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"elapsed":18047,"status":"ok","timestamp":1662041696085,"user":{"displayName":"Jacky Liang","userId":"05524594537381871813"},"user_tz":240},"id":"UCAim-bCO6ml","outputId":"1e5f7d10-2183-42f5-d647-faaa5cadf940"},"outputs":[{"output_type":"stream","name":"stdout","text":["Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n","Collecting openai\n","  Downloading openai-0.23.0.tar.gz (43 kB)\n","\u001b[K     |████████████████████████████████| 43 kB 1.5 MB/s \n","\u001b[?25h  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n","  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n","    Preparing wheel metadata ... \u001b[?25l\u001b[?25hdone\n","Collecting ratelimiter\n","  Downloading ratelimiter-1.2.0.post0-py3-none-any.whl (6.6 kB)\n","Requirement already satisfied: pandas>=1.2.3 in /usr/local/lib/python3.7/dist-packages (from openai) (1.3.5)\n","Requirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from openai) (4.1.1)\n","Collecting pandas-stubs>=1.1.0.11\n","  Downloading pandas_stubs-1.2.0.62-py3-none-any.whl (163 kB)\n","\u001b[K     |████████████████████████████████| 163 kB 8.9 MB/s \n","\u001b[?25hRequirement already satisfied: tqdm in /usr/local/lib/python3.7/dist-packages (from openai) (4.64.0)\n","Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from openai) (1.21.6)\n","Requirement already satisfied: openpyxl>=3.0.7 in /usr/local/lib/python3.7/dist-packages (from openai) (3.0.10)\n","Requirement already satisfied: requests>=2.20 in /usr/local/lib/python3.7/dist-packages (from openai) (2.23.0)\n","Requirement already satisfied: et-xmlfile in /usr/local/lib/python3.7/dist-packages (from openpyxl>=3.0.7->openai) (1.1.0)\n","Requirement already satisfied: pytz>=2017.3 in /usr/local/lib/python3.7/dist-packages (from pandas>=1.2.3->openai) (2022.2.1)\n","Requirement already satisfied: python-dateutil>=2.7.3 in /usr/local/lib/python3.7/dist-packages (from pandas>=1.2.3->openai) (2.8.2)\n","Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.7/dist-packages (from python-dateutil>=2.7.3->pandas>=1.2.3->openai) (1.15.0)\n","Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests>=2.20->openai) (2022.6.15)\n","Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests>=2.20->openai) (2.10)\n","Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests>=2.20->openai) (1.24.3)\n","Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests>=2.20->openai) (3.0.4)\n","Building wheels for collected packages: openai\n","  Building wheel for openai (PEP 517) ... \u001b[?25l\u001b[?25hdone\n","  Created wheel for openai: filename=openai-0.23.0-py3-none-any.whl size=54478 sha256=63136fdcb615dbe5cc5cd20fb166b85d6113871c7632a10c9b866d5e8db5e8c6\n","  Stored in directory: /root/.cache/pip/wheels/70/d5/31/f9f67660319d89e4f54501d27b1e90f88a3309c42ea4fd734c\n","Successfully built openai\n","Installing collected packages: pandas-stubs, ratelimiter, openai\n","Successfully installed openai-0.23.0 pandas-stubs-1.2.0.62 ratelimiter-1.2.0.post0\n"]}],"source":["! pip install openai ratelimiter\n","import openai\n","openai.api_key = openai_api_key"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"xktpxJza0MYS"},"outputs":[],"source":["from google.colab import output\n","output.enable_custom_widget_manager()"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"xdQki-sb6rk1"},"outputs":[],"source":["from copy import copy\n","from time import sleep\n","from tqdm.auto import trange, tqdm\n","\n","import numpy as np\n","\n","import ast\n","import astunparse\n","\n","from pygments import highlight\n","from pygments.lexers import PythonLexer\n","from pygments.formatters import TerminalFormatter\n","\n","from time import time\n","from ratelimiter import RateLimiter\n","\n","def limited_cb(until):\n","    duration = int(round(until - time()))\n","    print('Rate limited, sleeping for {:d} seconds'.format(duration))\n","\n","openai_rate_limiter = RateLimiter(max_calls=15, period=60, callback=limited_cb)\n","\n","def exec_safe(code_str, gvars, lvars):\n","  banned_phrases = ['import', '__']\n","  for phrase in banned_phrases:\n","    assert phrase not in code_str\n","  \n","  empty_fn = lambda *args, **kwargs: None\n","  custom_gvars = merge_dicts([\n","      gvars,\n","      {'exec': empty_fn, 'eval': empty_fn}\n","  ])\n","  exec(code_str, custom_gvars, lvars)\n","\n","default_query_kwargs = {\n","    'engine': model_name,\n","    'max_tokens': 512,\n","    'temperature': 0,\n","    'frequency_penalty': 0,\n","    'logprobs': 1\n","}\n","\n","lmp_cache = {}\n","\n","def lmp(base_prompt, query, stop_tokens=None, log=False, return_response=False, strip=False, use_cache=False, query_kwargs=None):\n","    prompt = f'{base_prompt}\\n{query}'\n","    \n","    if not use_cache or prompt not in lmp_cache:\n","      use_query_kwargs = copy(default_query_kwargs)\n","      if query_kwargs is not None:\n","        use_query_kwargs.update(query_kwargs)\n","      with openai_rate_limiter:\n","        while True:\n","          try:\n","            result = openai.Completion.create(\n","                prompt=prompt, stop=stop_tokens, **use_query_kwargs\n","            )['choices'][0]\n","            \n","            response = result['text']\n","            logp = np.mean(result['logprobs']['token_logprobs'])\n","            break\n","          except Exception as e:\n","            print('got err')\n","            print(e)\n","            print('retrying after 10')\n","            sleep(10)\n","            continue\n","\n","      if strip:\n","        response = response.strip()\n","\n","      lmp_cache[prompt] = {\n","          'response': response,\n","          'mean_log_prob': logp\n","      }\n","\n","    response = lmp_cache[prompt]['response']\n","\n","    if log:\n","      print(query)\n","      print(response)\n","\n","    if return_response:\n","      return response\n","\n","def lmp_fgen(prompt, f_name, f_sig, stop_tokens=['# define function:', '# example:'], recurse=False, \n","            recurse_level=0, max_recurse_level=4, use_cache=False,\n","             context_vars=None, bug_fix=False, log=False, return_src=False, query_kwargs=None, strip=False, info=''):\n","    query = f'# define function: {f_sig}.'\n","    if info:\n","      query = f'{query}\\n# info: {info}.'\n","    if query_kwargs is not None:\n","      query_kwargs = copy(query_kwargs)\n","      query_kwargs['temperature'] = 0\n","    f_src = lmp(prompt, query, stop_tokens=stop_tokens, log=False, return_response=True, use_cache=use_cache, query_kwargs=query_kwargs)\n","    if bug_fix:\n","        with openai_rate_limiter:\n","          f_src = openai.Edit.create(\n","            model='code-davinci-edit-001',\n","            input='# ' + f_src,\n","            temperature=0,\n","            instruction=\"Fix syntax errors. Keep same inputs and outputs. Only small changes. No comments.\",\n","          )['choices'][0]['text']\n","\n","    if strip:\n","      f_src = f_src.strip()\n","\n","    if context_vars is None:\n","        context_vars = {}\n","    gvars = context_vars\n","    lvars = {}\n","\n","    f_success = True\n","    try:\n","      exec_safe(f_src, gvars, lvars)\n","      f = lvars[f_name]\n","    except Exception as e:\n","      # print('error', f_sig)\n","      # print(e)\n","      # print(f_src)\n","      # print()\n","      f = lambda *args, **kargs: None   \n","      f_success = False \n","\n","    all_child_fs, all_child_f_srcs = {}, {}\n","    if recurse and recurse_level < max_recurse_level and f_success:\n","      # recursively define child_fs in the function body if needed\n","      f_def_body = None\n","      for node in ast.parse(f_src).body:\n","        if isinstance(node, ast.FunctionDef):\n","          f_def_body = astunparse.unparse(node.body)\n","      if f_def_body is not None:      \n","        potential_child_fs, potential_child_f_sigs = {}, {}\n","        f_parser = FunctionParser(potential_child_fs, potential_child_f_sigs)\n","        f_parser.visit(ast.parse(f_def_body))\n","        for potential_child_f_name, potential_child_f_sig in potential_child_f_sigs.items():\n","          if potential_child_f_name in potential_child_fs:\n","            potential_child_fs[potential_child_f_name] = potential_child_f_sig\n","\n","        for child_f_name, child_f_sig in potential_child_fs.items():\n","          all_vars = merge_dicts([context_vars, all_child_fs, lvars])\n","          if not var_exists(child_f_name, all_vars):\n","            child_fs, child_f_srcs = lmp_fgen(\n","                prompt, child_f_name, child_f_sig, \n","                stop_tokens=stop_tokens, \n","                context_vars=all_vars, \n","                bug_fix=bug_fix,\n","                log=False, \n","                recurse=True,\n","                recurse_level=recurse_level+1,\n","                return_src=True,\n","                use_cache=use_cache,\n","                query_kwargs=query_kwargs\n","              )\n","\n","            all_child_fs.update(child_fs)\n","            all_child_f_srcs.update(child_f_srcs)\n","\n","        if len(all_child_fs) > 0:\n","          # redefine parent f so newly created all_child_fs are in scope\n","          gvars = merge_dicts([context_vars, all_child_fs])\n","          lvars = {}\n","        \n","          exec_safe(f_src, gvars, lvars)\n","          \n","          f = lvars[f_name]\n","\n","    if log:\n","        to_print = highlight(f'{query}\\n{f_src}', PythonLexer(), TerminalFormatter())\n","        print(f'LMP FGEN created:\\n\\n{to_print}\\n')\n","\n","    fs = {\n","        f_name: f\n","    }\n","    fs.update(all_child_fs)\n","\n","    if return_src:\n","        f_srcs = {\n","            f_name: f_src\n","        }\n","        f_srcs.update(all_child_f_srcs)\n","\n","        return fs, f_srcs\n","    return fs\n","\n","def lmp_batch(base_prompt, cmds, stop_tokens=['# define'], strip=False, batch_size=20, query_kwargs=None, ret_logprobs=False, use_cache=False):\n","    prompts = [\n","      f'{base_prompt}\\n{cmd}'\n","      for cmd in cmds\n","    ]\n","\n","    if use_cache:\n","      prompts_use_idxs = [\n","          idx for idx, prompt in enumerate(prompts) if prompt not in lmp_cache\n","      ]\n","    else:\n","      prompts_use_idxs = list(range(len(prompts)))\n","\n","    use_query_kwargs = copy(default_query_kwargs)\n","    if query_kwargs is not None:\n","      use_query_kwargs.update(query_kwargs)\n","\n","    for start_idx in trange(0, len(prompts_use_idxs), batch_size, leave=False):\n","        end_idx = min(start_idx + batch_size, len(prompts_use_idxs))\n","        batch_idxs = prompts_use_idxs[start_idx : end_idx]\n","        batch_prompts = [prompts[idx] for idx in batch_idxs]\n","\n","        with openai_rate_limiter:\n","          while True:\n","            try:\n","              raw_responses_batch = openai.Completion.create(\n","                  prompt=batch_prompts, stop=stop_tokens, **use_query_kwargs\n","              )\n","              break\n","            except Exception as e:\n","              print('got err')\n","              print(e)\n","              print('retrying after 10')\n","              sleep(10)\n","              continue\n","            \n","        responses_batch = [\n","            r['text']\n","            for r in raw_responses_batch['choices']\n","        ]\n","        mean_logprobs_batch = [\n","            np.mean(r['logprobs']['token_logprobs'])\n","            for r in raw_responses_batch['choices']\n","        ]\n","\n","        if strip:\n","            responses_batch = [response.strip() for response in responses_batch]\n","\n","        for p, r, logp in zip(batch_prompts, responses_batch, mean_logprobs_batch):\n","          lmp_cache[p] = {\n","              'response': r,\n","              'mean_log_prob': logp\n","          }\n","\n","    responses = [lmp_cache[p]['response'] for p in prompts]\n","\n","    if ret_logprobs:\n","      mean_log_probs = [lmp_cache[p]['mean_log_prob'] for p in prompts]\n","      return responses, mean_log_probs\n","\n","    return responses\n","\n","def lmp_fgen_batch(prompt, prompt_with_comment, queries, stop_tokens=['# define function:', '# example:'], \n","                   recurse=False, context_vars=None, log=False, strip=False, query_kwargs=None, ret_logprobs=False):\n","\n","    f_srcs_list = lmp_batch(prompt, queries, stop_tokens=stop_tokens, query_kwargs=query_kwargs, ret_logprobs=ret_logprobs, use_cache=False)\n","    if ret_logprobs:\n","      f_srcs_list, logprobs = f_srcs_list\n","    for idx, (query, f_src) in enumerate(zip(queries, f_srcs_list)):\n","      f_srcs_list[idx] = query + f_src\n","    \n","    if strip:\n","      for idx, f_src in enumerate(f_srcs_list):\n","        f_srcs_list[idx] = f_src.strip()\n","\n","    if recurse:\n","      if context_vars is None:\n","        context_vars = {}\n","\n","      # recursively define child_fs in the function body if needed\n","      for idx, f_src in enumerate(f_srcs_list):\n","        try:\n","          lvars = {}\n","          exec(f_src, {}, lvars)\n","\n","          f_def_body = None\n","          for node in ast.parse(f_src).body:\n","            if isinstance(node, ast.FunctionDef):\n","              f_def_body = astunparse.unparse(node.body)\n","          assert f_def_body is not None\n","        except Exception as e:\n","          # print('err recurse')\n","          # print(e)\n","          # print(f_src)\n","          # print()\n","          continue\n","\n","        potential_child_fs, potential_child_f_sigs = {}, {}\n","        f_parser = FunctionParser(potential_child_fs, potential_child_f_sigs)\n","        f_parser.visit(ast.parse(f_def_body))\n","        for potential_child_f_name, potential_child_f_sig in potential_child_f_sigs.items():\n","          if potential_child_f_name in potential_child_fs:\n","            potential_child_fs[potential_child_f_name] = potential_child_f_sig\n","\n","        all_child_fs, all_child_f_srcs = {}, {}\n","        for child_f_name, child_f_sig in potential_child_fs.items():\n","          all_vars = merge_dicts([context_vars, all_child_fs, lvars])\n","          if not var_exists(child_f_name, all_vars):\n","            child_fs, child_f_srcs = lmp_fgen(\n","                prompt_with_comment, child_f_name, child_f_sig,\n","                context_vars=all_vars, \n","                log=False,\n","                recurse=True,\n","                return_src=True,\n","                use_cache=True,\n","                query_kwargs=query_kwargs\n","              )\n","\n","            all_child_fs.update(child_fs)\n","            all_child_f_srcs.update(child_f_srcs)\n","\n","        if len(all_child_fs) > 0:\n","          child_f_srcs_str = \"\\n\".join(all_child_f_srcs.values())\n","          f_srcs_list[idx] = f'{f_src}\\n{child_f_srcs_str}'\n","          \n","    if log:\n","      for query, f_src in zip(queries, f_srcs_list):\n","        to_print = highlight(f_src, PythonLexer(), TerminalFormatter())\n","        print(f'LMP FGEN created:\\n\\n{to_print}\\n')\n","\n","    if ret_logprobs:\n","      return f_srcs_list, logprobs\n","\n","    return f_srcs_list\n","\n","class FunctionParser(ast.NodeTransformer):\n","\n","    def __init__(self, fs, f_assigns):\n","      super().__init__()\n","      self._fs = fs\n","      self._f_assigns = f_assigns\n","\n","    def visit_Call(self, node):\n","        self.generic_visit(node)\n","        if isinstance(node.func, ast.Name):\n","            f_sig = astunparse.unparse(node).strip()\n","            f_name = astunparse.unparse(node.func).strip()\n","            self._fs[f_name] = f_sig\n","        return node\n","\n","    def visit_Assign(self, node):\n","        self.generic_visit(node)\n","        if isinstance(node.value, ast.Call):\n","            assign_str = astunparse.unparse(node).strip()\n","            f_name = astunparse.unparse(node.value.func).strip()\n","            self._f_assigns[f_name] = assign_str\n","        return node\n","\n","def var_exists(name, all_vars):\n","    try:\n","        eval(name, all_vars)\n","    except:\n","        exists = False\n","    else:\n","        exists = True\n","    return exists\n","\n","def merge_dicts(dicts):\n","    return {\n","        k : v \n","        for d in dicts\n","        for k, v in d.items()\n","    }"]},{"cell_type":"markdown","metadata":{"id":"BkSVrVQtO4BG"},"source":["# Run Benchmark"]},{"cell_type":"markdown","metadata":{"id":"9kreNge45P2G"},"source":["## Prompts"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"rnNe2xMSYDeM"},"outputs":[],"source":["prompt_f_gen_hier = '''\n","def get_total(xs: List[float]) -> float:\n","    \"\"\"Find the sum of a list of numbers called xs.\n","    \"\"\"\n","    return sum(xs)\n","# end of function\n","\n","def get_abs_diff_between_means(xs0: List[float], xs1: List[float]) -> float:\n","    \"\"\"Get the absolute difference between the means of two lists of numbers.\n","    \"\"\"\n","    m0 = get_mean(xs0)\n","    m1 = get_mean(xs1)\n","    return abs(m0 - m1)\n","# end of function\n","'''.strip()\n","\n","\n","prompt_f_gen_hier_comment = '''\n","# define function: total = get_total(xs).\n","def get_total(xs):\n","    return sum(xs)\n","\n","# define function: diff = get_abs_diff_between_means(xs0, xs1).\n","def get_abs_diff_between_means(xs0, xs1):\n","    m0 = get_mean_pure_python(xs0)\n","    m1 = get_mean_pure_python(xs1)\n","    return abs(m0 - m1)\n","'''.strip()\n","\n","prompt_f_gen_flat = '''\n","def get_total(xs: List[float]) -> float:\n","    \"\"\"Find the sum of a list of numbers called xs.\n","    \"\"\"\n","    return sum(xs)\n","\n","def get_abs_diff_between_means(xs0: List[float], xs1: List[float]) -> float:\n","    \"\"\"Get the absolute difference between the means of two lists of numbers.\n","    \"\"\"\n","    m0 = sum(xs0) / len(xs0)\n","    m1 = sum(xs1) / len(xs1)\n","    return abs(m0 - m1)\n","'''.strip()"]},{"cell_type":"markdown","metadata":{"id":"a5Nznh2b5RTK"},"source":["## Load Problems"]},{"cell_type":"code","execution_count":null,"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"elapsed":130,"status":"ok","timestamp":1662041696442,"user":{"displayName":"Jacky Liang","userId":"05524594537381871813"},"user_tz":240},"id":"jhT5hr_IMrAI","outputId":"54eadfb0-2079-453c-eeb1-154efe289612"},"outputs":[{"output_type":"stream","name":"stdout","text":["164\n","dict_keys(['task_id', 'prompt', 'entry_point', 'canonical_solution', 'test'])\n"]}],"source":["from human_eval.data import write_jsonl, read_problems\n","\n","problems = read_problems()\n","\n","task_ids = list(problems.keys())\n","task_prompts = [\n","                f\"{problem['prompt']}\"\n","                for problem in problems.values()\n","              ]\n","\n","print(len(problems))\n","print(problems['HumanEval/0'].keys())"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"1_rZBZLA5CuB"},"outputs":[],"source":["idx = 0\n","problem = problems[f'HumanEval/{idx}']\n","solutions = lmp_fgen_batch(prompt_f_gen_hier, prompt_f_gen_hier_comment, [problem['prompt']], stop_tokens=['def', 'if __name__', '# end of function'], recurse=True, log=True)"]},{"cell_type":"markdown","metadata":{"id":"5rTQn4a0EE8B"},"source":["## Hier greedy"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"9C-tF1LZ9XOM"},"outputs":[],"source":["solutions = lmp_fgen_batch(prompt_f_gen_hier, prompt_f_gen_hier_comment, task_prompts, stop_tokens=['def', 'if __name__'], recurse=True)\n","results = [\n","    {\n","        'task_id': task_id,\n","        'completion': solution\n","    }\n","    for task_id, solution in zip(task_ids, solutions)\n","]\n","result_path = results_path / \"results_hier_greedy.jsonl\"\n","write_jsonl(result_path, results)\n","\n","! python3 human-eval/human_eval/evaluate_functional_correctness.py \"$result_path\""]},{"cell_type":"markdown","metadata":{"id":"jMoshTL9EIEh"},"source":["## Flat w/ Prompt greedy"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"lUDxkCJBQXB6"},"outputs":[],"source":["solutions = lmp_fgen_batch(prompt_f_gen_flat, '', task_prompts, stop_tokens=['def', 'if __name__'], recurse=False)\n","results = [\n","    {\n","        'task_id': task_id,\n","        'completion': solution\n","    }\n","    for task_id, solution in zip(task_ids, solutions)\n","]\n","\n","result_path = results_path / \"results_flat_with_prompt_greedy.jsonl\"\n","write_jsonl(result_path, results)\n","\n","! python3 human-eval/human_eval/evaluate_functional_correctness.py \"$result_path\""]},{"cell_type":"markdown","metadata":{"id":"cSWqonHYEd5z"},"source":["## Flat w/o Prompt greedy"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"4H88WQWgEoFe"},"outputs":[],"source":["solutions = lmp_fgen_batch('', '', task_prompts, stop_tokens=['def', 'if __name__'], recurse=False)\n","results = [\n","    {\n","        'task_id': task_id,\n","        'completion': solution\n","    }\n","    for task_id, solution in zip(task_ids, solutions)\n","]\n","\n","result_path = results_path / \"results_flat_without_prompt_greedy.jsonl\"\n","\n","write_jsonl(result_path, results)\n","\n","! python3 human-eval/human_eval/evaluate_functional_correctness.py \"$result_path\""]},{"cell_type":"markdown","metadata":{"id":"agQ7OGflznaP"},"source":["## Hier samples"]},{"cell_type":"code","source":["results = []\n","for _ in trange(100):\n","  solutions, logprobs = lmp_fgen_batch(prompt_f_gen_hier, prompt_f_gen_hier_comment, task_prompts, stop_tokens=['def', 'if __name__', '# end of function'], recurse=True, query_kwargs={'temperature': 0.8}, ret_logprobs=True)\n","  results.extend([\n","      {\n","          'task_id': task_id,\n","          'completion': solution,\n","          'logprob': logprob\n","      }\n","      for task_id, solution, logprob in zip(task_ids, solutions, logprobs)\n","  ])\n","\n","result_path = results_path / \"results_hier_samples.jsonl\"\n","\n","write_jsonl(result_path, results)\n","\n","! python3 human-eval/human_eval/evaluate_functional_correctness.py \"$result_path\""],"metadata":{"id":"N45gOyIjZ0Dh"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"gPvRhMkXzqFL"},"source":["## Flat w/ Prompt samples"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"sRatxvm2zsjY"},"outputs":[],"source":["results = []\n","for _ in trange(100):\n","  solutions, logprobs = lmp_fgen_batch(prompt_f_gen_flat, '', task_prompts, stop_tokens=['def', 'if __name__'], recurse=False, query_kwargs={'temperature': 0.8}, ret_logprobs=True)\n","  results.extend([\n","      {\n","          'task_id': task_id,\n","          'completion': solution,\n","          'logprob': logprob\n","      }\n","      for task_id, solution, logprob in zip(task_ids, solutions, logprobs)\n","  ])\n","\n","result_path = results_path / \"results_flat_with_prompt_samples.jsonl\"\n","\n","write_jsonl(result_path, results)\n","\n","! python3 human-eval/human_eval/evaluate_functional_correctness.py \"$result_path\""]},{"cell_type":"markdown","metadata":{"id":"KOlGac32zsrr"},"source":["## Flat w/o Prompt samples"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"q0lH1fdrzvxP"},"outputs":[],"source":["results = []\n","for _ in trange(100):\n","  solutions, logprobs = lmp_fgen_batch('', '', task_prompts, stop_tokens=['def', 'if __name__'], recurse=False, query_kwargs={'temperature': 0.8}, ret_logprobs=True)\n","  results.extend([\n","      {\n","          'task_id': task_id,\n","          'completion': solution,\n","          'logprob': logprob\n","      }\n","      for task_id, solution, logprob in zip(task_ids, solutions, logprobs)\n","  ])\n","\n","result_path = results_path / \"results_flat_without_prompt_samples.jsonl\"\n","\n","write_jsonl(result_path, results)\n","\n","! python3 human-eval/human_eval/evaluate_functional_correctness.py \"$result_path\""]}],"metadata":{"colab":{"collapsed_sections":[],"provenance":[{"file_id":"1kQnLoe9BeXDHfuSltcE3EJ7aPMjjeYY7","timestamp":1658507602813}],"toc_visible":true},"kernelspec":{"display_name":"Python 3","name":"python3"},"language_info":{"name":"python"}},"nbformat":4,"nbformat_minor":0}